package com.zyx.flinkdemo.stream.sink.hbasebatch;

import com.alibaba.fastjson.JSONObject;
import com.zyx.flinkdemo.pojo.HttpDataModel;
import lombok.extern.slf4j.Slf4j;
import org.apache.flink.configuration.Configuration;
import org.apache.flink.streaming.api.functions.sink.RichSinkFunction;
import org.apache.hadoop.hbase.*;
import org.apache.hadoop.hbase.client.*;
import org.apache.hadoop.hbase.util.Bytes;

import java.io.Serializable;
import java.util.ArrayList;
import java.util.HashMap;
import java.util.List;
import java.util.Map;


/**
 * @author Yaxi.Zhang
 * @since 2021/8/10 11:28
 * desc: Flink写入HBase的RickSink类
 */
@Slf4j
public class HbaseSink extends RichSinkFunction<List<HttpDataModel>> implements Serializable {

    private final String hbaseZookeeperHost;
    private final String hbaseZookeeperPort;

    private Connection conn;
    private Admin admin;

    public HbaseSink(String hbaseZookeeperHost, String hbaseZookeeperPort) {
        this.hbaseZookeeperHost = hbaseZookeeperHost;
        this.hbaseZookeeperPort = hbaseZookeeperPort;
    }

    @Override
    public void open(Configuration parameters) throws Exception {
        super.open(parameters);

        org.apache.hadoop.conf.Configuration configuration = HBaseConfiguration.create();
        configuration.set("hbase.zookeeper.property.clientPort", hbaseZookeeperPort);
        configuration.set("hbase.zookeeper.quorum", hbaseZookeeperHost);

        conn = ConnectionFactory.createConnection(configuration);
        admin = conn.getAdmin();
    }

    @Override
    public void invoke(List<HttpDataModel> datas, Context context) throws Exception {
        // 按 project:table 归纳
        Map<String, List<HttpDataModel>> map = new HashMap<String, List<HttpDataModel>>();
        for (HttpDataModel data : datas) {
            if (! map.containsKey(data.getFullTable())) {
                map.put(data.getFullTable(), new ArrayList<HttpDataModel>());
            }
            map.get(data.getFullTable()).add(data);
        }
        // 遍历 map
        for(Map.Entry<String, List<HttpDataModel>> entry : map.entrySet()){
            // 如果 表不存在，即创建
            createTable(entry.getKey());
            // 写数据
            List<Put> list = new ArrayList<Put>();
            for (HttpDataModel item : entry.getValue()) {
                Put put = new Put(Bytes.toBytes(String.valueOf(System.currentTimeMillis())));

                JSONObject object = JSONObject.parseObject(item.getData());
                for (String key: object.keySet()) {
                    put.addColumn("data".getBytes(), key.getBytes(), object.getString(key).getBytes());
                }
                list.add(put);
            }
            conn.getTable(TableName.valueOf(entry.getKey())).put(list);
        }
    }

    @Override
    public void close() throws Exception {
        super.close();
    }

    /**
     * 创建 HBase 表
     */
    private void createTable(String tableName) throws Exception {
        createNamespace(tableName.split(":")[0]);
        TableName table = TableName.valueOf(tableName);
        if (!admin.tableExists(table)) {
            TableDescriptor hTableDescriptor = TableDescriptorBuilder
                    .newBuilder(table)
                    .setColumnFamily(ColumnFamilyDescriptorBuilder.of("data"))
                    .build();
            admin.createTable(hTableDescriptor);
        }
    }

    /**
     * 创建命名空间
     */
    private void createNamespace(String namespace) throws Exception {
        try {
            admin.getNamespaceDescriptor(namespace);
        } catch (NamespaceNotFoundException e) {
            admin.createNamespace(NamespaceDescriptor.create(namespace).build());
        }
    }

}
